import torch
import torch.nn as nn
import torch.nn.functional as F
import pdb

class CrowdCounter(nn.Module):
    def __init__(self,gpus,model_name):
        super(CrowdCounter, self).__init__()        

        if model_name == 'AlexNet':
            from .SCC_Model.AlexNet import AlexNet as net        
        elif model_name == 'VGG':
            from .SCC_Model.VGG import VGG as net
        elif model_name == 'VGG_DECODER':
            from .SCC_Model.VGG_decoder import VGG_decoder as net
        elif model_name == 'MCNN':
            from .SCC_Model.MCNN import MCNN as net
        elif model_name == 'CSRNet':
            from .SCC_Model.CSRNet import CSRNet as net
        elif model_name == 'Res50':
            from .SCC_Model.Res50 import Res50 as net
        elif model_name == 'Res101':
            from .SCC_Model.Res101 import Res101 as net            
        elif model_name == 'Res101_SFCN':
            from .SCC_Model.Res101_SFCN import Res101_SFCN as net
        elif model_name == 'SCAR':
            from .SCAR.SCAR import SCAR as net
        elif model_name == 'Stackpool_Base':
            from .stackpool_models.base import base as net
        elif model_name == 'Stackpool_Deep':
            from .stackpool_models.deep import deep as net
        elif model_name == 'Stackpool_Wide':
            from .stackpool_models.wide import wide as net
        elif model_name == 'DANet':
            from .DANet_model.model import Fine as net

        if model_name in ['AlexNet', 'VGG', 'VGG_DECODER', 'MCNN', 'CSRNet', 'Res50',
                          'Res101','Res101_SFCN', 'SCAR', 'DANet']:
            self.CCN = net()
        elif model_name in ['Stackpool_Base', 'Stackpool_Deep', 'Stackpool_Wide']:
            pools = ['vpool','stackpool','mpool']
            pool = pools[0]
            self.CCN = net(pool)

        if len(gpus)>1:
            self.CCN = torch.nn.DataParallel(self.CCN, device_ids=gpus).cuda()
        else:
            self.CCN=self.CCN.cuda()
        self.loss_mse_fn = nn.MSELoss().cuda()
        
    @property
    def loss(self):
        return self.loss_mse
    
    def forward(self, img, gt_map):                               
        density_map = self.CCN(img)

        if img.shape[0]==1: # batch size == 1
            # 对于MCNN,如果输出全尺寸的密度度，由于放大4倍，所以可能存在output与gt不一样大的问题，
            # 需要裁剪一下
            density_map_sq = density_map.squeeze()
            gt_map_sq = gt_map.squeeze()
            if gt_map_sq.shape != density_map_sq.shape:
                h, w  = density_map_sq.shape
                gt_map_sq = gt_map_sq[:h, :w]
            self.loss_mse= self.build_loss(density_map_sq, gt_map_sq)

        else: # 当batch size大于1时 dataloader会自动裁剪图像到统一的大小
            self.loss_mse= self.build_loss(density_map.squeeze(), gt_map.squeeze())
        return density_map
    
    def build_loss(self, density_map, gt_data):
        loss_mse = self.loss_mse_fn(density_map, gt_data)  
        return loss_mse

    def test_forward(self, img):                               
        density_map = self.CCN(img)                    
        return density_map

